import torch
from torch import nn
import torch.nn.functional as F
import math
#词汇表->指定维度embedding
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, dim_model):
        super(TokenEmbedding, self).__init__()
        self.vocab_size = vocab_size
        self.dim_model = dim_model
        self.embedding = nn.Embedding(vocab_size, dim_model)
    def forward(self, x):
        return self.embedding(x) * math.sqrt(self.dim_model)

class PositionalEmbedding(nn.Module):
    def __init__(self, dim_model, seq_len):
        super(PositionalEmbedding, self).__init__()
        self.pe = torch.zeros(seq_len, dim_model)
        self.pe.requires_grad = False
        pos = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(dim=1) 
        _2i = torch.arange(0, dim_model, step=2).float()
        div_term = torch.exp(_2i * (-math.log(10000.0) / dim_model))
        self.pe[:,0::2] = torch.sin(pos * div_term)
        self.pe[:,1::2] = torch.cos(pos * div_term) 
        self.pe = self.pe.unsqueeze(0)
    
    def forward(self, x):
        _, seq_len = x.size()
        return self.pe[:,:seq_len,:]

class TransformerEmbedding(nn.Module):
    def __init__(self, vocab_size, dim_model, seq_len, dropout_p):
        super(TransformerEmbedding,self).__init__()
        self.token_emb = TokenEmbedding(vocab_size, dim_model)
        self.pos_emb = PositionalEmbedding(dim_model, seq_len)
        self.dropout = nn.Dropout(p=dropout_p)
    
    def forward(self, x):
        token_emb = self.token_emb(x)
        pos_emb = self.pos_emb(x)
        return self.dropout(token_emb + pos_emb)
    



         

